#!/bin/bash
# queuedb - Manage AutoGluon Assistant WebUI task queue database

# Exit on errors
set -Eeuo pipefail

# Parse command line arguments
COMMAND=""
CUSTOM_DB_PATH=""

# Function to show usage
usage() {
    echo "Usage: $0 [--db-path PATH] (create|reset|dump)"
    echo "  --db-path PATH  - Custom database path (optional)"
    echo "  create          - Create and initialize the task queue database"
    echo "  reset           - Clear all data from the tasks table"
    echo "  dump            - Display all tasks in the database"
    echo ""
    echo "If --db-path is not specified, defaults to: \$HOME/.autogluon_assistant/webui_queue.db"
}

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --db-path)
            CUSTOM_DB_PATH="$2"
            shift 2
            ;;
        create|reset|dump)
            COMMAND="$1"
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            echo "Unknown option: $1"
            usage
            exit 1
            ;;
    esac
done

# Check if command is provided
if [ -z "$COMMAND" ]; then
    usage
    exit 1
fi

# Set database path
if [ -n "$CUSTOM_DB_PATH" ]; then
    DB_PATH="$CUSTOM_DB_PATH"
    DB_DIR=$(dirname "$DB_PATH")
else
    # Default path
    DB_DIR="$HOME/.autogluon_assistant"
    DB_PATH="$DB_DIR/webui_queue.db"
fi

# Execute command
case $COMMAND in
    "create")
        # Create directory if it doesn't exist
        mkdir -p "$DB_DIR"
        
        # Check if database file exists
        if [ -f "$DB_PATH" ]; then
            # Check if tasks table exists
            TABLE_EXISTS=$(sqlite3 "$DB_PATH" "SELECT name FROM sqlite_master WHERE type='table' AND name='tasks';" 2>/dev/null || echo "")
            if [ -n "$TABLE_EXISTS" ]; then
                echo "Database and tasks table already exist at $DB_PATH"
                exit 0
            fi
        fi

        # Create and initialize the database
        sqlite3 "$DB_PATH" << 'EOF'
CREATE TABLE IF NOT EXISTS tasks (
    id INTEGER PRIMARY KEY AUTOINCREMENT,
    task_id TEXT UNIQUE NOT NULL,
    run_id TEXT,
    status TEXT NOT NULL CHECK (status IN ('queued', 'running', 'completed', 'cancelled')),
    command_json TEXT NOT NULL,
    credentials_json TEXT,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    started_at TIMESTAMP
);

CREATE INDEX IF NOT EXISTS idx_status ON tasks(status);
CREATE INDEX IF NOT EXISTS idx_created_at ON tasks(created_at);
EOF
        echo "Database created and initialized at $DB_PATH"
        ;;

    "reset")
        # Check if database exists
        if [ ! -f "$DB_PATH" ]; then
            echo "Error: database does not exist at $DB_PATH"
            exit 1
        fi
        
        # Clear all data from tasks table
        sqlite3 "$DB_PATH" "DELETE FROM tasks;"
        
        # Reset the autoincrement counter
        sqlite3 "$DB_PATH" "DELETE FROM sqlite_sequence WHERE name='tasks';"
        
        echo "All tasks cleared from database at $DB_PATH"
        ;;

    "dump")
        # If the database does not exist, then exit
        if [ ! -f "$DB_PATH" ]; then
            echo "Error: database does not exist at $DB_PATH"
            exit 1
        fi

        echo "Database: $DB_PATH"
        echo ""
        
        # Show summary first
        echo "=== Task Queue Summary ==="
        sqlite3 "$DB_PATH" << 'EOF'
SELECT 
    'Total tasks: ' || COUNT(*) as summary,
    'Queued: ' || SUM(CASE WHEN status = 'queued' THEN 1 ELSE 0 END),
    'Running: ' || SUM(CASE WHEN status = 'running' THEN 1 ELSE 0 END),
    'Completed: ' || SUM(CASE WHEN status = 'completed' THEN 1 ELSE 0 END),
    'Cancelled: ' || SUM(CASE WHEN status = 'cancelled' THEN 1 ELSE 0 END)
FROM tasks;
EOF
        
        echo -e "\n=== All Tasks ==="
        # Dump all data from the tasks table with better formatting
        sqlite3 -column -header "$DB_PATH" << 'EOF'
SELECT 
    id,
    substr(task_id, 1, 8) || '...' as task_id,
    CASE WHEN run_id IS NOT NULL THEN substr(run_id, 1, 8) || '...' ELSE 'N/A' END as run_id,
    status,
    datetime(created_at, 'localtime') as created_at,
    CASE WHEN started_at IS NOT NULL THEN datetime(started_at, 'localtime') ELSE 'N/A' END as started_at
FROM tasks
ORDER BY id DESC;
EOF
        ;;

    *)
        usage
        exit 1
        ;;
esac
